package com.c2b.algorithm.leetcode.base;

import java.util.Deque;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.Map;

/**
 * <a href="https://leetcode.cn/problems/count-complete-tree-nodes/description/?envType=study-plan-v2&envId=top-interview-150">完全二叉树的节点个数(Count Complete Tree Nodes)</a>
 * <p>给你一棵 完全二叉树 的根节点 root ，求出该树的节点个数。</p>
 * <p><a href="https://baike.baidu.com/item/%E5%AE%8C%E5%85%A8%E4%BA%8C%E5%8F%89%E6%A0%91/7773232?fr=aladdin">完全二叉树</a>的定义如下：在完全二叉树中，除了最底层节点可能没填满外，其余每层节点数都达到最大值，并且最下面一层的节点都集中在该层最左边的若干位置。若最底层为第 h 层，则该层包含 1~ 2^h 个节点。</p>
 * <p>
 * <b>示例：</b>
 * <pre>
 * 示例 1：
 *      输入：root = [1,2,3,4,5,6]
 *      输出：6
 *
 * 示例 2：
 *      输入：root = []
 *      输出：0
 *
 * 示例 3：
 *      输入：root = [1]
 *      输出：1
 *     </pre>
 * </p>
 * <p>
 * <b>提示：</b>
 * <ul>
 *     <li>树中节点的数目范围是[0, 5 * 10^4]</li>
 *     <li>0 <= Node.val <= 5 * 10^4</li>
 *     <li>题目数据保证输入的树是 完全二叉树</li>
 * </ul>
 * </p>
 * <b>进阶：遍历树来统计节点是一种时间复杂度为 O(n) 的简单解决方案。你可以设计一个更快的算法吗？</b>
 *
 * @author c2b
 * @since 2023/10/18 13:59
 */
public class LC0222CountCompleteTreeNodes_M {

    static class Solution {
        public int countNodes(TreeNode root) {
            if (root == null) {
                return 0;
            }
            if (root.left == null && root.right == null) {
                return 1;
            }
            int maxLevel = 0;
            // 节点所在层的映射
            Map<TreeNode, Integer> levelMap = new HashMap<>();
            Deque<TreeNode> stack = new LinkedList<>();
            while (root != null) {
                stack.push(root);
                maxLevel++;
                levelMap.put(root, maxLevel);
                root = root.left;
            }
            int lastLevelNodeCount = 0;
            while (!stack.isEmpty()) {
                TreeNode currNode = stack.poll();
                Integer currLevel = levelMap.get(currNode);
                if (currLevel == maxLevel) {
                    continue;
                }
                if (currLevel == maxLevel - 1) {
                    if (currNode.left != null) {
                        lastLevelNodeCount++;
                    }
                    if (currNode.right != null) {
                        lastLevelNodeCount++;
                    }
                } else {
                    currNode = currNode.right;
                    while (currNode != null) {
                        stack.push(currNode);
                        currLevel++;
                        levelMap.put(currNode, currLevel);
                        currNode = currNode.left;
                    }
                }
            }
            return (int) (Math.pow(2, maxLevel - 1) - 1 + lastLevelNodeCount);
        }
    }

    public static void main(String[] args) {
        Solution solution = new Solution();
        TreeNode root = new TreeNode(1);
        //root.left = new TreeNode(2);
        //root.right = new TreeNode(3);
        //root.left.left = new TreeNode(4);
        //root.left.right = new TreeNode(5);
        //root.right.left = new TreeNode(6);
        System.out.println(solution.countNodes(root));
    }
}
